/**
 * Copyright (c) 2018-2023, NXOS Development Team
 * SPDX-License-Identifier: Apache-2.0
 *
 * Contains: Implementation of virtio-net
 *
 * Change Logs:
 * Date           Author            Notes
 * 2023-08-08     Shiroko           Init
 */

#include "nxos.h"

#ifdef CONFIG_NX_DRIVER_VIRTIO_NET

#undef NX_LOG_MOD_NAME
#define NX_LOG_MOD_NAME "VIRTIO-NET"

#include "base/log.h"

#define DRV_NAME "virtio net"
#define DEV_NAME "virtio_net"

#include "../virtio.h"
#include "./virtio_net.h"

NX_PRIVATE NX_Error VirtioNetTransmit(NX_Device *device, NX_U8 *buffer, NX_Size size)
{
    if (size > VIRTIO_NET_MSS) return NX_EINVAL;
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    // Get far most idx from used
    VirtioQueue              *transmitq       = VirtioQueueGet(device, VIRTIO_NET_QUEUE_TRANSMIT);
    NX_U32                   idx              =
                                     (transmitq->avail->idx + transmitq->ring_size - 1) % transmitq->ring_size;
    VirtioNetHeader          *t_header        = &deviceExtension->transmit_buffer[idx].header;
    NX_U8                    *t_buffer        = deviceExtension->transmit_buffer[idx].buffer;
    // Temporarily the packet is simple one without additional features
    t_header->flags       = 0;
    t_header->gso_type    = VIRTIO_NET_HDR_GSO_NONE;
    t_header->num_buffers = 0;
    NX_MemCopy(t_buffer, buffer, size);
    VirtioQueueDescGet(transmitq, idx)->len = sizeof(VirtioNetHeader) + size;
    VirtioQueuePutIntoAvailRing(transmitq, idx);
    VirtioSendNotification(device, VIRTIO_NET_QUEUE_TRANSMIT, VIRTIO_QUEUE_DESC_INVALID_ID);
    return NX_EOK;
}

NX_PRIVATE NX_Error VirtioNetQueueReceiveHandler(NX_Device *device, VirtioQueue *queue, NX_U32 idx)
{
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    // VirtioNetHeader          *header          = &deviceExtension->receive_buffer[idx].header;
    // NX_U8                    *buffer          = deviceExtension->receive_buffer[idx].buffer;
    // NX_Size                  buffer_size      = queue->used->ring[idx].len - sizeof(VirtioNetHeader);
    // NX_Printf("\t\tReceived packet, length: %d bytes.\n", buffer_size);
    // PrintBuffer(buffer, buffer_size, 2);
    if (deviceExtension->received_ids != NX_NULL)
        NX_FifoWrite(deviceExtension->received_ids, (NX_U8 *) &idx, sizeof(NX_U32));
    VirtioQueuePutIntoAvailRing(queue, idx);
    VirtioSendNotification(device, VIRTIO_NET_QUEUE_RECEIVE, idx);
    return NX_EOK;
}


NX_PRIVATE void VirtioNetCleanUp(NX_Device *device);

NX_PRIVATE NX_Error VirtioNetInit(NX_Driver *driver, NX_Device *device)
{
    NX_Error err                       = NX_EOK;
    NX_U32   features[1];
    features[0] =
            VIRTIO_NET_FEATURE_MAC       // Device has a given mac address
            | VIRTIO_NET_FEATURE_STATUS;  // Status field in DevConfig is available
    // features[1] = VIRTIO_FEATURE_VERSION_1;
    // When the network stack is finished, we could enable all possible features.

    // Negotiate features & status
    if ((err = VirtioDeviceInitialize(device, features, 1)) != NX_EOK)
        return err;

    // Device-specified initialization
    VirtioQueueAlloc(device, 2);
    VirtioQueueInit(device, VIRTIO_NET_QUEUE_RECEIVE, CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);
    VirtioQueueInit(device, VIRTIO_NET_QUEUE_TRANSMIT, CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE);
    // We do not enable control queue
    VirtioQueueRegisterNotifier(VirtioQueueGet(device, VIRTIO_NET_QUEUE_RECEIVE),
                                VirtioNetQueueReceiveHandler);

    // Driver OK
    VirtioDeviceDriverOK(device);

    VirtioNetDeviceExtension *deviceExtension = NX_MemAlloc(sizeof(VirtioNetDeviceExtension));
    if (deviceExtension == NX_NULL) goto nomem_failed;
    VirtioGetExtension(device)->devExt = (void *) deviceExtension;

    VirtioNetDeviceConfig *deviceCfg = NX_MemAlloc(sizeof(VirtioNetDeviceConfig));
    if (deviceCfg == NX_NULL) goto nomem_failed;
    VirtioReadDeviceConfig(device, 0, deviceCfg, sizeof(VirtioNetDeviceConfig));

    NX_MemCopy(deviceExtension->mac, deviceCfg->mac, 6);
    char macStr[3 * 6 + 1];
    for (int i = 0; i < 6; i++)
        NX_SNPrintf(macStr + i * 3, 3, "%02x:", deviceExtension->mac[i]);
    macStr[5 * 3 + 2] = '\0';
    NX_LOG_I("Net interface MAC address: %s", macStr);

    /*
     * Set no interrupt for transmit queue for performance.
     * Using the most far descriptor from avail->idx for new transmission, that must already be used by the device.
     */
    VirtioQueue *transmitq = VirtioQueueGet(device, VIRTIO_NET_QUEUE_TRANSMIT);
    transmitq->avail->flags |= VIRTIO_QUEUE_AVAIL_FLAG_NO_INTERRUPT;
    for (int idx = 0; idx < CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE; idx++)
    {
        VirtioQueueDescFill(transmitq, idx,
                            (NX_U64) ((NX_UArch) NX_Virt2Phy(&deviceExtension->transmit_buffer[idx])),
                            sizeof(deviceExtension->transmit_buffer[idx]),
                            0, 0);
    }

    /*
     * Setup receive buffers
     */
    VirtioQueue *receiveq = VirtioQueueGet(device, VIRTIO_NET_QUEUE_RECEIVE);
    for (NX_U32 idx       = 0; idx < CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE; idx++)
    {
        NX_ASSERT(receiveq->free[idx] == NX_True);
        receiveq->free[idx] = NX_False;
        receiveq->free_count -= 1;
        VirtioQueueDescFill(receiveq, idx,
                            (NX_U64) ((NX_UArch) NX_Virt2Phy(&deviceExtension->receive_buffer[idx])),
                            sizeof(deviceExtension->receive_buffer[idx]),
                            VIRTIO_QUEUE_DESC_FLAG_WRITE, 0);
        VirtioQueuePutIntoAvailRing(receiveq, idx);
        receiveq->trace[idx].needNotify = NX_True;
    }
    VirtioSendNotification(device, VIRTIO_NET_QUEUE_RECEIVE, VIRTIO_QUEUE_DESC_INVALID_ID);

    return NX_EOK;
nomem_failed:
    err = NX_ENOMEM;

    VirtioNetCleanUp(device);
    return err;
}

NX_PRIVATE void VirtioNetCleanUp(NX_Device *device)
{
    // Reset the device
    VirtioDeviceReset(device);
    // Cleanup VirtioQueue
    for (int i = 0; i < VirtioGetExtension(device)->queue_num; i++)
        VirtioQueueDestroy(device, i);
    VirtioQueueFree(device);

    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    if (deviceExtension == NX_NULL) return;
    NX_MemFree(deviceExtension);
}

NX_PRIVATE NX_Error open(struct NX_Device *device, NX_U32 flags)
{
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    if (deviceExtension->received_ids != NX_NULL) return NX_EBUSY;
    // Obviously using fifo for network packet receiving will possibly drop many of them. This is for testing before we got real network functionality
    if ((deviceExtension->received_ids =
                 NX_FifoCreate(sizeof(NX_U32) * CONFIG_NX_DRIVER_VIRTIO_SPLIT_QUEUE_SIZE)) == NX_NULL)
        return NX_ENOMEM;
    return NX_EOK;
}

NX_PRIVATE NX_Error close(struct NX_Device *device)
{
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    NX_FifoDestroy(deviceExtension->received_ids);
    deviceExtension->received_ids = NX_NULL;
    return NX_EOK;
}

NX_PRIVATE NX_Error read(struct NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    if (deviceExtension->received_ids == NX_NULL) return NX_ENORES;
    NX_U32 idx = 0;
    if (NX_FifoRead(deviceExtension->received_ids, (NX_U8 *) &idx, sizeof(NX_U32)) != sizeof(NX_U32))
        return NX_ENORES;
    VirtioQueue *receiveq = VirtioQueueGet(device, VIRTIO_NET_QUEUE_RECEIVE);

    // VirtioNetHeader *header     = &deviceExtension->receive_buffer[idx].header;
    NX_U8           *buffer     = deviceExtension->receive_buffer[idx].buffer;
    NX_Size         buffer_size = receiveq->used->ring[idx].len - sizeof(VirtioNetHeader);

    if (len < buffer_size) return NX_ENOMEM;

    NX_MemCopy(VirtioVirt2Phy(buf), buffer, buffer_size);
    *outLen = buffer_size;

    return NX_EOK;
}

NX_PRIVATE NX_Error write(struct NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    NX_Error err;
    if ((err = VirtioNetTransmit(device, VirtioVirt2Phy(buf), len)) != NX_EOK)
        return err;
    *outLen = len;
    return NX_EOK;
}

NX_PRIVATE NX_Error control(struct NX_Device *device, NX_U32 cmd, void *arg)
{
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    switch (cmd)
    {
    case VIRTIO_NET_CMD_GET_MAC:
        NX_MemCopy(VirtioVirt2Phy(arg), deviceExtension->mac, sizeof(deviceExtension->mac));
        break;
    case VIRTIO_NET_CMD_FLUSH:
        if (deviceExtension->received_ids == NX_NULL) return NX_ENORES;
        NX_FifoReset(deviceExtension->received_ids);
        break;
    default:
        return NX_ENOFUNC;
    }
    return NX_EOK;
}

NX_PRIVATE NX_Error mappable(struct NX_Device *device, NX_Size length, NX_U32 port, NX_Addr *outPhyAddr)
{
    // Maybe we could map rx/tx buffer to userspace to reduce internal copy
    return NX_ENOFUNC;
}

NX_PRIVATE NX_Error poll(struct NX_Device *device, NX_PollState *pState)
{
    VirtioNetDeviceExtension *deviceExtension = VirtioGetDevExtension(device, VirtioNetDeviceExtension);
    if (deviceExtension->received_ids == NX_NULL) return NX_ENORES;
    return NX_FifoPoll(deviceExtension->received_ids, pState);
}


NX_PRIVATE NX_DriverOps virtioNetDriverOps = {
        .open = open,
        .close = close,
        .read = read,
        .write = write,
        .control = control,
        .mappable = mappable,
        .poll = poll
};


NX_PRIVATE VirtioDriverTableEntity virtioNetDriverTableEntity = {
        .deviceId = VIRTIO_DEVICE_ID_NET,
        .driverName = "virtio-net",
        .deviceNameTemplate = "virtio_net_%d",
        .deviceType = NX_DEVICE_TYPE_NET,
        .driverOps = &virtioNetDriverOps,
        .initHandler = VirtioNetInit,
        .cleanUpHandler = VirtioNetCleanUp
};

VirtioDriverTableEntity *VirtioDriverTableEntityForVirtioNet(void)
{
    return &virtioNetDriverTableEntity;
}

#else

#include "../virtio.h"

VirtioDriverTableEntity *VirtioDriverTableEntityForVirtioNet(void)
{
    return NX_NULL;
}

#endif // CONFIG_NX_DRIVER_VIRTIO_NET
